import json
import sys
import psutil
from tqdm import tqdm
from s2c.engine.system_info import SystemInfo
from s2c.engine.process_info import ProcessInfo, field_names
from s2c.rules.conf import FIELD_DELIMITER, PROCESS_TAB,SYS_BASE,SYS_CURRENT
from s2c.engine.db_helper import DbHelper


# ProcessScanner 进程表扫描器
# dict_type: current | base， current表示用户扫描当前系统，base 基础指纹库
class ProcessScanner:
    def __init__(self, dict_type=SYS_CURRENT):
        self.sys = None
        self.ps_list = []
        self.dict_type = dict_type
        self.pids = []

    # 直接使用psutil库方式获取当前系统所有进程
    def scan(self):
        self.pids = psutil.pids()
        total = len(self.pids)
        bar = tqdm(total=total, desc='进程扫描:', unit='', unit_scale=False)
        for pid in self.pids:
            info = ProcessInfo(pid, dict_type=self.dict_type)
            info.detail()
            self.ps_list.append(info.fields())
            bar.update(1)
        print("进程扫描完成！ 总数: " + str(len(self.ps_list)))

    def process_data(self):
        print('-----------', self.ps_list)
        # 如果是基础镜像，不对数据库进行操作
        if self.dict_type == SYS_BASE:
            return
        db = DbHelper()
        new_arr = []
        
        for item in self.ps_list:
            new = db.check_db_system_pro_table(item)
            if new:
                new_arr.append(item)
        print('\n---------扫描的新进程信息---------\n')
        for item in new_arr:
            print(item['user'] + '	' + item['name'])
        return new_arr

    # 将default字典写入到JSON文件中
    def save(self):
        self.sys = SystemInfo()
        # 先获取当前系统default 信息，并生成文件名称
        obj = self.sys.obj()
        file_name = self.sys.process_name(dict_type=self.dict_type)
        obj['table_name'] = PROCESS_TAB
        obj['fields'] = field_names()
        obj['delimiter'] = FIELD_DELIMITER
        obj['lines'] = len(self.ps_list)  # 增加数据的条数记录到对象中
        obj['dict_type'] = self.dict_type
        data_arr = {
            'info': obj,
            'list': self.ps_list
        }
        with open(file_name, 'w') as f: json.dump(data_arr, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    scanner = ProcessScanner(SYS_BASE)
    scanner.scan()
    scanner.save()
